//=====================================================================================
// 1)  This module implements the round coefficients
//     The round i corresponds to the polynomial x^(i - 1) in GF(2^8) for i = 1,...,10
//=====================================================================================
module round_coeff
(
  //OUTPUTS
  output reg [7:0] r_coeff,
  //INPUTS
  input [3:0] round
);
  always @(*)
    begin
      case(round)
        4'd0 : r_coeff = 8'h01;
        4'd1 : r_coeff = 8'h02;
        4'd2 : r_coeff = 8'h04;
        4'd3 : r_coeff = 8'h08;
        4'd4 : r_coeff = 8'h10;
        4'd5 : r_coeff = 8'h20;
        4'd6 : r_coeff = 8'h40;
        4'd7 : r_coeff = 8'h80;
        4'd8 : r_coeff = 8'h1b;
        4'd9 : r_coeff = 8'h36;
        default: r_coeff = 8'h00;
      endcase
    end
endmodule

//=====================================================================================
// 2)  This module implements the counter round convertion necessary to decryption
//=====================================================================================
module round_count_convertion
(
  //OUTPUTS
  output reg [3:0] count_conv,
  //INPUTS
  input [3:0] count
);
    
  always @(*)
    begin
      case(count)
        4'd0 : count_conv = 4'd9;
        4'd1 : count_conv = 4'd8;
        4'd2 : count_conv = 4'd7;
        4'd3 : count_conv = 4'd6;
        4'd4 : count_conv = 4'd5;
        4'd5 : count_conv = 4'd4;
        4'd6 : count_conv = 4'd3;
        4'd7 : count_conv = 4'd2;
        4'd8 : count_conv = 4'd1;
        4'd9 : count_conv = 4'd0;
        default: count_conv = 4'd0;
      endcase
    end

endmodule

//=====================================================================================
// 3)  This module implements the Rotation Word Operation
//=====================================================================================
module rot_word
#(
  parameter DATA_WIDTH = 32,
  parameter WORD_SIZE = 8
)(
  //OUTPUTS
  output [DATA_WIDTH - 1:0] word_rot,
  //INPUTS
  input [DATA_WIDTH - 1:0] word
);
  generate
    genvar i;
    for(i = 0; i < DATA_WIDTH; i = i + 1)
      begin
        assign word_rot[i] = word[(i + WORD_SIZE)%DATA_WIDTH];
      end
  endgenerate
  
endmodule

//=====================================================================================
// 4)  This module implements the g function
//=====================================================================================
module g_function
(
  //OUTPUTS
  output [31:0] g_out,
  //INPUTS
  input [31:0] g_in,
  input [ 3:0] round,
  input enc_dec
);

  wire [31:0] word_rot;
  wire [ 7:0] rc;
  wire [ 7:0] sbox0_out;
  wire [ 7:0] sbox1_out;
  wire [ 7:0] sbox2_out;
  wire [ 7:0] sbox3_out;
  
  rot_word ROTW
  (
    .word_rot(word_rot),
    .word(g_in)
  );
  
  sBox_8 SBOX_0
  (
    .sbox_out ( sbox0_out     ),
    .sbox_in  ( word_rot[7:0] ),
    .enc_dec  ( enc_dec       )
  );
  
    sBox_8 SBOX_1
  (
    .sbox_out ( sbox1_out      ),
    .sbox_in  ( word_rot[15:8] ),
    .enc_dec  ( enc_dec        )
  );
  
    sBox_8 SBOX_2
  (
    .sbox_out ( sbox2_out       ),
    .sbox_in  ( word_rot[23:16] ),
    .enc_dec  ( enc_dec         )
  );
  
    sBox_8 SBOX_3
  (
    .sbox_out ( sbox3_out       ),
    .sbox_in  ( word_rot[31:24] ),
    .enc_dec  ( enc_dec         )
  );
  
  round_coeff RC
  (
    .r_coeff ( rc    ),
    .round   ( round )
  );
  
  assign g_out = {sbox3_out, sbox2_out, sbox1_out, sbox0_out ^ rc};
  
endmodule

//===============================================================================================
// 5)  This module implements the combinational logic of round key expander with dedicated S-Box
//===============================================================================================
module key_expander_sbox
(
  //OUTPUTS
  output [31:0] word0_out,
  output [31:0] word1_out,
  output [31:0] word2_out,
  output [31:0] word3_out,
  //INPUTS
  input [31:0] word0,
  input [31:0] word1,
  input [31:0] word2,
  input [31:0] word3,
  input [ 3:0]round,
  input enc_dec
);

  wire [31:0] g_out;

  g_function G_FUNC
  (
    .g_out   ( g_out   ),
    .g_in    ( word3   ),
    .round   ( round   ),
    .enc_dec ( enc_dec )
  );

  assign word0_out = word0 ^ g_out; 
  assign word1_out = word1 ^ word0_out;
  assign word2_out = word2 ^ word1_out;
  assign word3_out = word3 ^ word2_out;
endmodule

//===============================================================================================
// 6)  This module implements the combinational logic of round key expander with shared S-Box
//===============================================================================================
module key_expander
(
  //OUTPUTS
  output [31:0] key0_out,     //Key 0 for round i + 1
  output [31:0] key1_out,     //Key 1 for round i + 1
  output [31:0] key2_out,     //Key 2 for round i + 1
  output [31:0] key3_out,     //Key 3 for round i + 1
  output [31:0] rot_word,     //Key 3 after rotation word operation
  //INPUTS
  input [31:0] key0,          //Key 0 for round i
  input [31:0] key1,          //Key 1 for round i
  input [31:0] key2,          //Key 2 for round i
  input [31:0] key3,          //Key 3 for round i
  input [31:0] sbox,          //Key 3 after rotation word and S-Box operations
  input [ 3:0] round,         //Number of current round (0 to 10)
  input enc_dec               //Encryption/Decryption selection          
);

  wire [ 3:0] round_conv;
  wire [ 3:0] rc_in;
  wire [ 7:0] rc_out;
  wire [31:0] g_out;
  
  rot_word ROTW
  (
    .word_rot ( rot_word ),
    .word     ( key3     )
  );

  round_count_convertion COUNT_CONV
  (
    .count_conv ( round_conv),
    .count      ( round )
  );
  
  assign rc_in = enc_dec ? round : round_conv; 
  
  round_coeff RC
  (
    .r_coeff ( rc_out ),
    .round   ( rc_in  )
  );
  
  assign g_out = {sbox[31:24], sbox[23:16], sbox[15:8], sbox[7:0] ^ rc_out};
  
  assign key0_out = key0 ^ g_out; 
  assign key1_out = key1 ^ key0;
  assign key2_out = key2 ^ key1;
  assign key3_out = key3 ^ key2;
  
endmodule
